from typing import Dict, Any
from scorpio.core.database import VectorDatabaseManager
from .document import DocumentProcessor
from .query_history import QueryHistoryProcessor

from scorpio.core.common import get_logger,EmbeddingConfig
logger = get_logger(__name__)

class DatabaseInitializerService:
    def __init__(self, db_manager: VectorDatabaseManager, embedding_config: EmbeddingConfig):
        logger.debug("initialize database initializer service")
        self.db_manager = db_manager
        self.embedding_config = embedding_config
        self.document_processor = DocumentProcessor(db_manager, embedding_config.dimension)
        self.query_history_processor = QueryHistoryProcessor(db_manager, embedding_config.dimension)
        self.initialized = False
    
    async def _check_is_initialized(self) -> bool:
        """检查是否所有必要的表都已存在"""
        async with self.db_manager.get_connection() as conn:
            existing_tables = await conn.fetch("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = 'public' 
                    AND table_name IN ('documents', 'chunks', 'query_history');
            """)
        
        existing_table_names = {table["table_name"] for table in existing_tables}
        required_tables = {"documents", "chunks", "query_history"}
        
        if existing_table_names == required_tables:
            logger.info("all required tables has been initialized.")
            return True
        
        # 记录缺失的表，便于调试
        missing_tables = required_tables - existing_table_names
        if missing_tables:
            logger.info(f"检测到缺失的表: {missing_tables}")
        
        return False
    
    async def initialize(self) -> bool:
        """完整的数据库初始化流程"""
        if self.initialized:
            logger.info("Database has been initialized, skip to initialize.")
            return True
        try:
            logger.info("start to initialize database...")
            # 步骤1: 检查数据库连接
            await self._test_connection()
            
            # 步骤2: 检查是否所有必要的表都已存在
            if await self._check_is_initialized():
                self.initialized = True
                return True
            
            # 后续初始化步骤...
            await self._create_extensions()
            
            # 步骤3: 创建表结构
            await self._create_tables()
            
            # 步骤4: 创建索引
            await self._create_indexes()
            
            # 步骤5: 验证初始化结果
            health_status = await self.check_database_health()
            
            if health_status["status"] == "healthy":
                self.initialized = True
                logger.info("database has been initialized.")
                return True
            else:
                logger.error("database initialization verification failed.")
                return False
                
        except Exception as e:
            logger.error(f"database initialization failed: {e}")
            raise
    
    async def _test_connection(self):
        """测试数据库连接"""
        try:
            # 通过 VectorDatabaseManager 测试连接
            async with self.db_manager.get_connection() as conn:
                await conn.execute("SELECT 1")
            logger.info("database connection test passed.")
        except Exception as e:
            logger.error(f"database connection test failed: {e}")
            raise
    
    async def _create_extensions(self):
        """创建必要的 PostgreSQL 扩展"""
        async with self.db_manager.get_connection() as conn:
            # 启用 vector 扩展
            await conn.execute('CREATE EXTENSION IF NOT EXISTS vector;')
            logger.info("vector extension has been enabled.")
            
            # 启用其他可能需要的扩展
            await conn.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp";')
            logger.info("uuid-ossp extension has been enabled.")
    
    async def _create_tables(self):
        """创建数据库表结构"""
        # 使用 DocumentProcessor 类来创建document相关的表
        await self.document_processor.create_table()
        
        # 使用 QueryHistoryProcessor 类来创建查询历史相关的表
        await self.query_history_processor.create_table()
        
    
    async def _create_indexes(self):
        """创建性能索引"""
        await self.document_processor.create_indexes()
        await self.query_history_processor.create_indexes()
    
    async def check_database_health(self) -> Dict[str, Any]:
        """检查数据库健康状态"""
        try:
            async with self.db_manager.get_connection() as conn:
                # 检查表是否存在
                tables_check = await conn.fetch('''
                    SELECT table_name 
                    FROM information_schema.tables 
                    WHERE table_schema = 'public' 
                    AND table_name IN ('documents', 'chunks', 'query_history');
                ''')
                
                existing_tables = [row['table_name'] for row in tables_check]
                
                # 检查向量扩展
                vector_ext = await conn.fetchval(
                    "SELECT COUNT(*) FROM pg_extension WHERE extname = 'vector';"
                )
                
                # 获取表统计信息
                doc_count = await conn.fetchval("SELECT COUNT(*) FROM documents;")
                chunk_count = await conn.fetchval("SELECT COUNT(*) FROM chunks;")
                
                return {
                    "status": "healthy",
                    "vector_extension": vector_ext > 0,
                    "existing_tables": existing_tables,
                    "document_count": doc_count,
                    "chunk_count": chunk_count,
                    "required_tables": ["documents", "chunks", "query_history"],
                    "all_tables_exist": len(existing_tables) == 3
                }
                
        except Exception as e:
            return {
                "status": "unhealthy",
                "error": str(e)
            }
    
    async def reset(self, application_env: str):
        """重置数据库（仅用于开发和测试）"""
        if application_env == "production":
            raise ValueError("生产环境禁止重置数据库")
            
        async with self.db_manager.get_connection() as conn:
            # 删除表（注意顺序，因为有外键约束）
            await conn.execute('DROP TABLE IF EXISTS query_history CASCADE;')
            await conn.execute('DROP TABLE IF EXISTS chunks CASCADE;')
            await conn.execute('DROP TABLE IF EXISTS documents CASCADE;')
            
            logger.warning("所有表已删除")
            
            # 重新初始化
            self.initialized = False
            await self.initialize()
    
    async def get_database_schema_info(self) -> Dict[str, Any]:
        """获取数据库模式信息"""
        async with self.db_manager.get_connection() as conn:
            # 获取表结构信息
            tables_info = await conn.fetch('''
                SELECT 
                    table_name,
                    column_name,
                    data_type,
                    is_nullable,
                    column_default
                FROM information_schema.columns 
                WHERE table_schema = 'public'
                ORDER BY table_name, ordinal_position;
            ''')
            
            # 获取索引信息
            indexes_info = await conn.fetch('''
                SELECT 
                    indexname,
                    tablename,
                    indexdef
                FROM pg_indexes 
                WHERE schemaname = 'public'
                ORDER BY tablename, indexname;
            ''')
            
            return {
                "tables": [
                    {
                        "table_name": row["table_name"],
                        "column_name": row["column_name"],
                        "data_type": row["data_type"],
                        "is_nullable": row["is_nullable"],
                        "column_default": row["column_default"]
                    }
                    for row in tables_info
                ],
                "indexes": [
                    {
                        "index_name": row["indexname"],
                        "table_name": row["tablename"],
                        "definition": row["indexdef"]
                    }
                    for row in indexes_info
                ]
            }
    async def get_database_stats(self) -> Dict[str, Any]:
        """获取数据库统计信息"""
        async with self.db_manager.get_connection() as conn:
            # 基本统计
            stats = {
                "total_documents": await conn.fetchval("SELECT COUNT(*) FROM documents;"),
                "total_chunks": await conn.fetchval("SELECT COUNT(*) FROM chunks;"),
                "total_queries": await conn.fetchval("SELECT COUNT(*) FROM query_history;")
            }
            
            # 文档类型分布
            content_types = await conn.fetch('''
                SELECT content_type, COUNT(*) as count 
                FROM documents 
                GROUP BY content_type;
            ''')
            stats["content_types"] = {
                row["content_type"]: row["count"] for row in content_types
            }
            
            # 向量维度信息
            if stats["total_chunks"] > 0:
                stats["vector_dimensions"] = await conn.fetchval('''
                    SELECT vector_dims(embedding) 
                    FROM chunks 
                    LIMIT 1;
                ''')
            else:
                stats["vector_dimensions"] = None
            
            return stats
    